{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hMMdZ1vnxj8f"
   },
   "source": [
    "<center>\n",
    "    <p style=\"text-align:center\">\n",
    "        <img alt=\"phoenix logo\" src=\"https://storage.googleapis.com/arize-phoenix-assets/assets/phoenix-logo-light.svg\" width=\"200\"/>\n",
    "        <br>\n",
    "        <a href=\"https://arize.com/docs/phoenix/\">Docs</a>\n",
    "        |\n",
    "        <a href=\"https://github.com/Arize-ai/phoenix\">GitHub</a>\n",
    "        |\n",
    "        <a href=\"https://arize-ai.slack.com/join/shared_invite/zt-2w57bhem8-hq24MB6u7yE_ZF_ilOYSBw#/shared-invite/email\">Community</a>\n",
    "    </p>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "se3M7GECJfpZ"
   },
   "source": [
    "# AutoGen Routing Agent\n",
    "\n",
    "In this tutorial, we'll explore Routing agent workflows using [AutoGen's function calling capabilities and agent interactions](https://microsoft.github.io/autogen-for-net/articles/Create-an-agent.html) to direct tasks efficiently.\n",
    "\n",
    "Agent Routing is a pattern designed to handle incoming requests by classifying them and directing them to the single most appropriate specialized agent or workflow. This ensures requests are immediately handled by the correct specialist, improving efficiency.\n",
    "\n",
    "AutoGen simplifies implementing this pattern by enabling a dedicated 'Router Agent' to analyze incoming messages and use function calls to signal its classification decision. Based on this classification, the workflow explicitly directs the query to the appropriate specialist agent for a focused, separate interaction.\n",
    "\n",
    "With Phoenix tracing, you get full visibility into this process, clearly seeing the router's classification, the function call made, and the subsequent handoff to the specialist.\n",
    "\n",
    "By the end of this tutorial, you’ll learn how to:\n",
    "\n",
    "- Define AutoGen agents with specific tools (functions).\n",
    "- Create a dedicated Router Agent for request classification.\n",
    "- Implement a workflow that routes requests to specialist agents based on classification.\n",
    "- Trace and visualize the routing process and agent interactions using Phoenix.\n",
    "\n",
    "⚠️ You'll need an OpenAI Key for this tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "G2j6_RvFJfpa"
   },
   "source": [
    "## Set up Keys and Dependencies\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qqq arize-phoenix arize-phoenix-otel openinference-instrumentation-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qq pyautogen==0.9 autogen-agentchat~=0.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "import autogen\n",
    "\n",
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass(\"🔑 Enter your OpenAI API key: \")\n",
    "\n",
    "if \"PHOENIX_API_KEY\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_API_KEY\"] = getpass(\"🔑 Enter your Phoenix API key: \")\n",
    "\n",
    "if \"PHOENIX_COLLECTOR_ENDPOINT\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_COLLECTOR_ENDPOINT\"] = getpass(\"🔑 Enter your Phoenix Collector Endpoint\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ibmucjY_QkaH"
   },
   "source": [
    "## Configure Tracing\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from phoenix.otel import register\n",
    "\n",
    "tracer_provider = register(\n",
    "    project_name=\"autogen-agents\",\n",
    "    auto_instrument=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kTVRWf0NJfpc"
   },
   "source": [
    "## Example Routing Task: Customer Support Agent\n",
    "\n",
    "We will build an intelligent customer service system, designed to efficiently handle diverse user queries by leveraging specialized agents. We use tracing to tackle the common challenge of ensuring incoming requests are promptly routed to the appropriate agent.\n",
    "\n",
    "Below is a high-level diagram of our setup. Specifically, a dedicated router agent analyzes the user's initial message solely to classify its topic, triggering a handoff. A specialized agent then uses its tools to address the query. Once a termination condition is reached, the user gets an output.\n",
    "\n",
    "![Diagram](https://storage.googleapis.com/arize-phoenix-assets/assets/images/autogen_routing_diagram.png)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qq5TvGKzO685"
   },
   "source": [
    "## Define Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dpjxM6BGLGW0"
   },
   "source": [
    "The `llm_config` specifies the configuration used for each `AssistantAgent`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_list = {\n",
    "    \"model\": \"gpt-4o\",\n",
    "    \"api_key\": os.environ[\"OPENAI_API_KEY\"],\n",
    "}\n",
    "llm_config = {\"config_list\": config_list}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W-zcucoH-kZD"
   },
   "source": [
    "This cell defines functions that simulate typical operations for a customer support scenario, such as retrieving user subscription data, checking payment history, verifying server status, and handling account actions like password resets. In a real setting, we would use data retrieval techniques.\n",
    "\n",
    "These functions will act as the tools for our agents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function Tools\n",
    "import json\n",
    "\n",
    "\n",
    "def lookup_subscription(user_id: str):\n",
    "    print(f\"lookup_subscription for user: {user_id}\")\n",
    "    subscriptions = {\n",
    "        \"user123\": {\"plan\": \"Pro\", \"status\": \"Active\", \"renewal_date\": \"2025-12-01\"},\n",
    "        \"user456\": {\"plan\": \"Free\", \"status\": \"Active\", \"renewal_date\": None},\n",
    "    }\n",
    "    data = subscriptions.get(user_id)\n",
    "    return json.dumps(data)\n",
    "\n",
    "\n",
    "def check_payment_history(user_id: str, limit: int = 3):\n",
    "    print(f\"check_payment_history for user: {user_id}\")\n",
    "    history = {\n",
    "        \"user123\": [\n",
    "            {\"date\": \"2025-04-01\", \"amount\": 20.00, \"status\": \"Paid\"},\n",
    "            {\"date\": \"2025-03-01\", \"amount\": 20.00, \"status\": \"Paid\"},\n",
    "        ],\n",
    "        \"user456\": [],\n",
    "    }\n",
    "    data = history.get(user_id, [])\n",
    "    return json.dumps(data[:limit])\n",
    "\n",
    "\n",
    "def check_server_status(service_name: str):\n",
    "    status = {\n",
    "        \"login_service\": \"Operational\",\n",
    "        \"payment_gateway\": \"Operational\",\n",
    "        \"app_server\": \"Not Operational\",\n",
    "    }\n",
    "    data = {\"service\": service_name, \"status\": status.get(service_name, \"Unknown\")}\n",
    "    return json.dumps(data)\n",
    "\n",
    "\n",
    "def request_password_reset(user_id: str):\n",
    "    print(f\"request_password_reset for user: {user_id}\")\n",
    "    success = user_id in [\"user123\", \"user456\"]\n",
    "    data = {\"user_id\": user_id, \"reset_initiated\": success}\n",
    "    return json.dumps(data)\n",
    "\n",
    "\n",
    "def check_account_status(user_id: str):\n",
    "    status = {\n",
    "        \"user123\": \"Active\",\n",
    "        \"user456\": \"Active\",\n",
    "        \"user789\": \"Inactive\",\n",
    "    }\n",
    "    data = {\"user_id\": user_id, \"status\": status.get(user_id, \"Not Found\")}\n",
    "    return json.dumps(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u_Ue1b2l_LsU"
   },
   "source": [
    "Next, we define three AutoGen `AssistantAgent` instances, each specialized for a specific customer support function: Billing, Technical support, and Account Management.\n",
    "\n",
    "Each agent is given a unique system message describing its role, and its `llm_config` includes schemas for only the specific function-calling tools relevant to its domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Billing Agent\n",
    "billing_agent = autogen.AssistantAgent(\n",
    "    name=\"Billing_Support_Agent\",\n",
    "    system_message=\"You are a billing support specialist. You can help with subscriptions, payments, and refunds. Use your available tools.\",\n",
    "    llm_config={\n",
    "        \"config_list\": config_list,\n",
    "        \"functions\": [\n",
    "            {\n",
    "                \"name\": \"lookup_subscription\",\n",
    "                \"description\": \"Get user subscription details.\",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\"user_id\": {\"type\": \"string\", \"description\": \"The user ID\"}},\n",
    "                    \"required\": [\"user_id\"],\n",
    "                },\n",
    "            },\n",
    "            {\n",
    "                \"name\": \"check_payment_history\",\n",
    "                \"description\": \"Check recent payment history.\",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"user_id\": {\"type\": \"string\", \"description\": \"The user ID\"},\n",
    "                        \"limit\": {\"type\": \"integer\", \"description\": \"Max number of records\"},\n",
    "                    },\n",
    "                    \"required\": [\"user_id\"],\n",
    "                },\n",
    "            },\n",
    "        ],\n",
    "    },\n",
    ")\n",
    "\n",
    "# Technical Agent\n",
    "technical_agent = autogen.AssistantAgent(\n",
    "    name=\"Technical_Support_Agent\",\n",
    "    system_message=\"You are a technical support specialist. You can help troubleshoot issues, check service status, and search the knowledge base. Use your available tools.\",\n",
    "    llm_config={\n",
    "        \"config_list\": config_list,\n",
    "        \"functions\": [\n",
    "            {\n",
    "                \"name\": \"check_server_status\",\n",
    "                \"description\": \"Check the status of a backend service.\",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"service_name\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"Name of the service (e.g., login_service)\",\n",
    "                        }\n",
    "                    },\n",
    "                    \"required\": [\"service_name\"],\n",
    "                },\n",
    "            },\n",
    "        ],\n",
    "    },\n",
    ")\n",
    "\n",
    "# Account Management Agent\n",
    "account_agent = autogen.AssistantAgent(\n",
    "    name=\"Account_Management_Agent\",\n",
    "    system_message=\"You are an account management specialist. You can help with account status and password resets. Use your available tools.\",\n",
    "    llm_config={\n",
    "        \"config_list\": config_list,\n",
    "        \"functions\": [\n",
    "            {\n",
    "                \"name\": \"request_password_reset\",\n",
    "                \"description\": \"Initiate a password reset for the user.\",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\"user_id\": {\"type\": \"string\", \"description\": \"The user ID\"}},\n",
    "                    \"required\": [\"user_id\"],\n",
    "                },\n",
    "            },\n",
    "            {\n",
    "                \"name\": \"check_account_status\",\n",
    "                \"description\": \"Check if a user account is active.\",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\"user_id\": {\"type\": \"string\", \"description\": \"The user ID\"}},\n",
    "                    \"required\": [\"user_id\"],\n",
    "                },\n",
    "            },\n",
    "        ],\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OxhOrKO0_thE"
   },
   "source": [
    "This section defines two helper functions:\n",
    "- `route_to_specialist` gets the topic classification decision made by a router LLM into a global variable\n",
    "- `check_for_termination_message` identifies if a message signals the end of the chat.\n",
    "\n",
    "Then, we initialize the main AutoGen `UserProxyAgent` and configuring it to use these functions for termination checks and equipping it to execute any of the tool functions when called by an `AssisstantAgent`. The `UserProxyAgent` also acts as the human user within the AutoGen framework, initiating conversations and managing user interaction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def route_to_specialist(topic: str, user_query_for_specialist: str):\n",
    "    print(f\"Router determined topic: {topic}\")\n",
    "    global routing_decision\n",
    "    routing_decision = {\"topic\": topic, \"query\": user_query_for_specialist}\n",
    "\n",
    "    if topic in [\"billing\", \"technical\", \"account\"]:\n",
    "        return f\"Routing request about '{topic}'\"\n",
    "    else:\n",
    "        return \"Could not determine a category.\"\n",
    "\n",
    "\n",
    "def check_for_termination_message(message_dict):\n",
    "    content = message_dict.get(\"content\")\n",
    "    if isinstance(content, str):\n",
    "        return content.rstrip().endswith(\"exit\")\n",
    "    return False\n",
    "\n",
    "\n",
    "user_proxy = autogen.UserProxyAgent(\n",
    "    name=\"User_Proxy\",\n",
    "    system_message=\"A human user. You will execute functions when requested.\",\n",
    "    human_input_mode=\"TERMINATE\",\n",
    "    is_termination_msg=check_for_termination_message,\n",
    "    function_map={\n",
    "        \"lookup_subscription\": lookup_subscription,\n",
    "        \"check_payment_history\": check_payment_history,\n",
    "        \"check_server_status\": check_server_status,\n",
    "        \"request_password_reset\": request_password_reset,\n",
    "        \"check_account_status\": check_account_status,\n",
    "        \"route_to_specialist\": route_to_specialist,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yUxBFORMAQDu"
   },
   "source": [
    "Finally, we define our router. This `AssistantAgent` classifies incoming customer support requests into categories: billing, technical, account, or unknown.\n",
    "\n",
    "The system message instructs the `router_agent` not to answer queries directly, but instead to analyze the request and call the specific `route_to_specialist` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "router_agent = autogen.AssistantAgent(\n",
    "    name=\"Support_Router\",\n",
    "    system_message=\"\"\"You are a customer support request router. Your job is to analyze the user's request and determine the primary topic: 'billing', 'technical', or 'account'.\n",
    "    Based on the primary topic, call the 'route_to_specialist' function with the identified topic and the user query.\n",
    "    Do NOT attempt to answer the user's query yourself. Only classify and call the function.\n",
    "    If the query is ambiguous, very general, or doesn't fit these categories, use the topic 'unknown'.\n",
    "    \"\"\",\n",
    "    llm_config={\n",
    "        \"config_list\": config_list,\n",
    "        \"temperature\": 0,\n",
    "        \"functions\": [\n",
    "            {\n",
    "                \"name\": \"route_to_specialist\",\n",
    "                \"description\": \"Route the user's request to the appropriate specialist based on the primary topic.\",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"topic\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"The topic identified ('billing', 'technical', 'account', or 'unknown').\",\n",
    "                        },\n",
    "                        \"user_query_for_specialist\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"The original user query to be passed to the specialist.\",\n",
    "                        },\n",
    "                    },\n",
    "                    \"required\": [\"topic\", \"user_query_for_specialist\"],\n",
    "                },\n",
    "            }\n",
    "        ],\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6Ed8ui7-_fW2"
   },
   "source": [
    "## Run Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B6-6u97rDTzp"
   },
   "source": [
    "We are now ready to run our agent. First, the user will be prompted for their query. Then there is an interaction between the `UserProxyAgent` and the `router_agent` to classify the request's topic.\n",
    "\n",
    "Based on the classification result captured from the first interaction, it either initiates a separate, follow-up chat with the appropriate `AssisstantAgent` or provides feedback to the user if routing was unsuccessful."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "routing_decision = None\n",
    "specialist_map = {\n",
    "    \"billing\": billing_agent,\n",
    "    \"technical\": technical_agent,\n",
    "    \"account\": account_agent,\n",
    "}\n",
    "\n",
    "initial_query = input(\"Please enter your customer support question: \")\n",
    "\n",
    "tracer = tracer_provider.get_tracer(__name__)\n",
    "with tracer.start_as_current_span(\n",
    "    \"CustomerSupportInteraction\",\n",
    "    openinference_span_kind=\"agent\",\n",
    ") as agent_span:\n",
    "    user_proxy.initiate_chat(\n",
    "        router_agent,\n",
    "        message=initial_query,\n",
    "        max_turns=2,\n",
    "    )\n",
    "\n",
    "    print(f\"Routing decision: {routing_decision}\")\n",
    "\n",
    "    selected_specialist = None\n",
    "    query_for_specialist = None\n",
    "\n",
    "    if routing_decision and routing_decision[\"topic\"] in specialist_map:\n",
    "        selected_specialist = specialist_map[routing_decision[\"topic\"]]\n",
    "        query_for_specialist = routing_decision[\"query\"]\n",
    "        print(f\"Routing to: {selected_specialist.name}\")\n",
    "\n",
    "        user_proxy.initiate_chat(\n",
    "            selected_specialist,\n",
    "            message=f\"Hello {selected_specialist.name}, I need help with the following: {query_for_specialist}\",\n",
    "            max_turns=5,\n",
    "        )\n",
    "\n",
    "    elif routing_decision:\n",
    "        print(\n",
    "            f\"Routing determined topic: {routing_decision['topic']}. No specific specialist assigned.\"\n",
    "        )\n",
    "        print(\n",
    "            \"Can you please specify if your issue relates to billing, technical problems, or account management?\"\n",
    "        )\n",
    "    else:\n",
    "        print(\"Routing failed.\")\n",
    "\n",
    "# Example Queries To Test:\n",
    "# initial_query = \"Hi, I think I was double charged last month for user user123.\"\n",
    "# initial_query = \"The login page isn't loading, is the login_service down?\"\n",
    "# initial_query = \"I forgot my password for user456, can you help?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pwbMnN9EfGWF"
   },
   "source": [
    "## View Results in Phoenix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WyS3R8I1fv0u"
   },
   "source": [
    "In Phoenix, we unlock observability into our agent's operations and decision-making process. The traces allow easy inspection of foundational elements like system prompts and function call. Phoenix also showcases the reasoning behind routing choices and provides transparency into tool interactions by tracking both input parameters and generated outputs.\n",
    "\n",
    "Run the cell below to see the full tracing results.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "\n",
    "HTML(\"\"\"\n",
    "<video width=\"800\" height=\"600\" controls>\n",
    "  <source src=\"https://storage.googleapis.com/arize-phoenix-assets/assets/gifs/autogen_router_results.mp4\" type=\"video/mp4\">\n",
    "  Your browser does not support the video tag.\n",
    "</video>\n",
    "\"\"\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
